import sys
import re
import os
import requests
import gzip
from tqdm import tqdm
import hashlib
from dataclasses import dataclass
import logging
from datetime import datetime
import argparse
import sqlite3
from typing import List, Optional

# 配置 logging
script_dir = os.path.dirname(os.path.abspath(__file__))
localsrc_dir = os.path.join(script_dir, 'localsrc')
# 确保 localsrc_dir 目录存在
os.makedirs(localsrc_dir, exist_ok=True)
log_dir = os.path.join(localsrc_dir, 'logs')
os.makedirs(log_dir, exist_ok=True)
# 设置日志编码为 UTF-8
logging.basicConfig(filename=os.path.join(log_dir, 'debug.log'), level=logging.DEBUG,
                    format='%(asctime)s - %(levelname)s - %(message)s', encoding='utf-8')


@dataclass
class PackageInfo:
    package: str
    version: str
    architecture: str
    maintainer: str
    installed_size: str
    depends: str
    suggests: str
    provides: str
    filename: str
    size: str
    md5sum: str
    sha1: str
    sha256: str
    section: str
    priority: str
    multi_arch: str
    homepage: str
    description: str
    original_maintainer: str


@dataclass
class RepositoryInfo:
    name: str
    url: str
    architecture: str
    dist: str
    source_types: list


class DBManager:
    def __init__(self, db_path: str):
        self.db_path = db_path
        self.conn = sqlite3.connect(db_path)
        self._create_table()

    def _create_table(self):
        cursor = self.conn.cursor()
        cursor.execute('''
            CREATE TABLE IF NOT EXISTS PackageInfo (
                package TEXT PRIMARY KEY,
                version TEXT,
                architecture TEXT,
                maintainer TEXT,
                installed_size TEXT,
                depends TEXT,
                suggests TEXT,
                provides TEXT,
                filename TEXT,
                size TEXT,
                md5sum TEXT,
                sha1 TEXT,
                sha256 TEXT,
                section TEXT,
                priority TEXT,
                multi_arch TEXT,
                homepage TEXT,
                description TEXT,
                original_maintainer TEXT
            )
        ''')
        self.conn.commit()

    def insert_package(self, package_info: PackageInfo):
        cursor = self.conn.cursor()
        cursor.execute('''
            INSERT OR REPLACE INTO PackageInfo VALUES (
                ?, ?, ?, ?, ?, ?, ?, ?, ?, ?,
                ?, ?, ?, ?, ?, ?, ?, ?, ?
            )
        ''', (
            package_info.package, package_info.version, package_info.architecture,
            package_info.maintainer, package_info.installed_size, package_info.depends,
            package_info.suggests, package_info.provides, package_info.filename,
            package_info.size, package_info.md5sum, package_info.sha1,
            package_info.sha256, package_info.section, package_info.priority,
            package_info.multi_arch, package_info.homepage, package_info.description,
            package_info.original_maintainer
        ))
        self.conn.commit()

    def get_package(self, package_name: str) -> Optional[PackageInfo]:
        cursor = self.conn.cursor()
        cursor.execute('SELECT * FROM PackageInfo WHERE package = ?', (package_name,))
        row = cursor.fetchone()
        if not row:
            return None
        return PackageInfo(*row)

    def search_packages(self, keyword: str) -> List[PackageInfo]:
        cursor = self.conn.cursor()
        cursor.execute('SELECT * FROM PackageInfo WHERE package LIKE ?', (f'%{keyword}%',))
        return [PackageInfo(*row) for row in cursor.fetchall()]

    def clear_all_packages(self):
        cursor = self.conn.cursor()
        cursor.execute('DELETE FROM PackageInfo')
        self.conn.commit()

    def close(self):
        self.conn.close()


class Repository:
    def __init__(self, repo_info):
        self.repo_info = repo_info
        self.db_manager = DBManager(os.path.join(localsrc_dir, f'package_cache_{repo_info.name}.db'))
        self.load_cache()

    def get_packages_data(self, source_type):
        url = f"{self.repo_info.url}/dists/{self.repo_info.dist}/{source_type}/binary-{self.repo_info.architecture}/Packages.gz"
        cache_file = os.path.join(localsrc_dir, 'dist', source_type,
                                  f"Packages_{source_type}_{self.repo_info.name}.gz")
        return self._get_cached_or_download(url, cache_file)

    def update_cache(self):
        logging.debug(f"开始更新 {self.repo_info.name} 软件包缓存...")
        print(f"正在更新 {self.repo_info.name} 软件包缓存...")
        
        # 先清空数据库
        self.db_manager.clear_all_packages()
        
        for source_type in self.repo_info.source_types:
            content = self.get_packages_data(source_type)
            if content:
                print(f"开始将 {source_type} 源的包信息缓存到 {self.repo_info.name} 的 package_cache.db...")
                self._parse_and_store_packages(content)
                print(f"{source_type} 源的包信息已成功缓存到 {self.repo_info.name} 的 package_cache.db。")
        logging.debug(f"{self.repo_info.name} 软件包缓存更新完成。")
        print(f"{self.repo_info.name} 软件包缓存更新完成。")

    def load_cache(self):
        cache_file = os.path.join(localsrc_dir, f'package_cache_{self.repo_info.name}.db')
        if os.path.exists(cache_file):
            try:
                logging.debug(f"成功从 {cache_file} 加载缓存数据。")
                print(f"成功从 {cache_file} 加载缓存数据。")
            except Exception as e:
                logging.error(f"加载 {cache_file} 时出错: {e}")
                print(f"加载 {cache_file} 时出错: {e}", file=sys.stderr)
                self.update_cache()
        else:
            self.update_cache()

    def _get_cached_or_download(self, url, cache_file):
        os.makedirs(os.path.dirname(cache_file), exist_ok=True)
        if os.path.exists(cache_file):
            try:
                logging.debug(f"从缓存文件 {cache_file} 获取信息...")
                print(f"从缓存文件 {cache_file} 获取信息...")
                with open(cache_file, 'rb') as f:
                    return gzip.decompress(f.read()).decode('utf-8')
            except Exception as e:
                logging.error(f"读取缓存文件 {cache_file} 时出错: {e}")
                print(f"读取缓存文件 {cache_file} 时出错: {e}", file=sys.stderr)
        return self._download_and_cache(url, cache_file)

    def _download_and_cache(self, url, cache_file):
        try:
            logging.debug(f"尝试从 {url} 获取信息并缓存...")
            print(f"开始下载 {url}...")
            response = requests.get(url, stream=True)
            total_size = int(response.headers.get('content-length', 0))
            block_size = 1024
            with open(cache_file, 'wb') as f, tqdm(
                    desc=f"下载 {os.path.basename(cache_file)}",
                    total=total_size,
                    unit='iB',
                    unit_scale=True,
                    unit_divisor=1024
            ) as bar:
                for data in response.iter_content(block_size):
                    bar.update(len(data))
                    f.write(data)
            logging.debug(f"成功从 {url} 获取信息并缓存。")
            print(f"成功下载并缓存 {url} 到 {cache_file}。")
            with open(cache_file, 'rb') as f:
                return gzip.decompress(f.read()).decode('utf-8')
        except requests.RequestException as e:
            logging.error(f"请求 {url} 时出错: {e}")
            print(f"请求 {url} 时出错: {e}", file=sys.stderr)
        except gzip.BadGzipFile:
            logging.error(f"无法解压缩 {url} 文件。")
            print(f"无法解压缩 {url} 文件。", file=sys.stderr)
        return None

    def _parse_and_store_packages(self, content):
        blocks = re.split(r'\n\n', content)
        total = len([b for b in blocks if b and re.search(r'Package: (.*)', b)])
        with tqdm(total=total, desc=f"导入包信息", unit='pkg') as pbar:
            for block in blocks:
                if block and re.search(r'Package: (.*)', block):
                    package_name = re.search(r'Package: (.*)', block).group(1)
                    info = self._create_package_info(block)
                    if info:
                        self.db_manager.insert_package(info)
                    pbar.update(1)

    def _create_package_info(self, package_block):
        get_match = lambda pattern: re.search(pattern, package_block).group(1) if re.search(pattern, package_block) else None
        return PackageInfo(
            package=get_match(r"Package: (.*)"),
            version=get_match(r"Version: (.*)"),
            architecture=get_match(r"Architecture: (.*)"),
            maintainer=get_match(r"Maintainer: (.*)"),
            installed_size=get_match(r"Installed-Size: (.*)"),
            depends=get_match(r"Depends: (.*)"),
            suggests=get_match(r"Suggests: (.*)"),
            provides=get_match(r"Provides: (.*)"),
            filename=get_match(r"Filename: (.*)"),
            size=get_match(r"Size: (.*)"),
            md5sum=get_match(r"MD5sum: (.*)"),
            sha1=get_match(r"SHA1: (.*)"),
            sha256=get_match(r"SHA256: (.*)"),
            section=get_match(r"Section: (.*)"),
            priority=get_match(r"Priority: (.*)"),
            multi_arch=get_match(r"Multi-Arch: (.*)"),
            homepage=get_match(r"Homepage: (.*)"),
            description=get_match(r"Description: (.*?)\n\n"),
            original_maintainer=get_match(r"Original-Maintainer: (.*)")
        )
        
    def get_package_info(self, package_name):
        package_name = self._process_package_name(package_name)
        return self.db_manager.get_package(package_name)

    def _process_package_name(self, package_name):
        if ':any' in package_name:
            package_name = package_name.replace(':any', '')
        elif '-' in package_name:
            parts = package_name.split('-')
            for i in range(len(parts), 0, -1):
                new_name = '-'.join(parts[:i])
                if self.db_manager.get_package(new_name):
                    return new_name
        return package_name

    def query_software_info(self, package_name, fuzzy_search=False):
        logging.debug(
            f"开始在 {self.repo_info.name} 中查询 {'包含 ' + package_name if fuzzy_search else package_name} 的软件包信息...")
        print(
            f"开始在 {self.repo_info.name} 中查询 {'包含 ' + package_name if fuzzy_search else package_name} 的软件包信息...")
        if fuzzy_search:
            matched_packages = self.db_manager.search_packages(package_name)
            if not matched_packages:
                logging.warning(f"在 {self.repo_info.name} 中未找到包含 {package_name} 的软件包信息。")
                print(f"在 {self.repo_info.name} 中未找到包含 {package_name} 的软件包信息。", file=sys.stderr)
                return
            if len(matched_packages) > 1:
                print("找到多个匹配的软件包，请选择要查看的软件包：")
                for i, pkg in enumerate(matched_packages, start=1):
                    print(f"{i}. {pkg.package}")
                print("输入 0 退出查询，输入单个序号查询对应包信息，输入多个序号（用空格、中英文逗号分割）或序号区间（用 - 号表示）查询多个包信息。")
                while True:
                    choice = input("请输入对应的编号：").strip()
                    if choice == '0':
                        logging.debug("用户选择退出查询。")
                        print("退出查询。")
                        return
                    selected_indices = self._parse_choice(choice, len(matched_packages))
                    if selected_indices:
                        for index in selected_indices:
                            info = matched_packages[index - 1]
                            self._print_package_info(info)
                        logging.debug("查询完成。")
                        print("查询完成。")
                        return
                    print("输入无效，请重新输入。")
            else:
                info = matched_packages[0]
                self._print_package_info(info)
                logging.debug("查询完成。")
                print("查询完成。")
        else:
            info = self.get_package_info(package_name)
            if info:
                self._print_package_info(info)
                logging.debug("查询完成。")
                print("查询完成。")
            else:
                logging.warning(f"在 {self.repo_info.name} 中未找到 {package_name} 的软件包信息。")
                print(f"在 {self.repo_info.name} 中未找到 {package_name} 的软件包信息。", file=sys.stderr)

    def _parse_choice(self, choice, max_index):
        selected_indices = set()
        for part in re.split(r'[ ,，]', choice):
            if '-' in part:
                try:
                    start, end = map(int, part.split('-'))
                    if 1 <= start <= end <= max_index:
                        selected_indices.update(range(start, end + 1))
                    else:
                        return []
                except ValueError:
                    return []
            else:
                try:
                    index = int(part)
                    if 1 <= index <= max_index:
                        selected_indices.add(index)
                    else:
                        return []
                except ValueError:
                    return []
        return sorted(selected_indices)

    def _print_package_info(self, package_info):
        print(f"{package_info.package}:{package_info.architecture} 版本信息：")
        print(f"  - {package_info.version}")
        print("\n包信息：")
        for attr in vars(package_info):
            value = getattr(package_info, attr)
            print(f"{attr.capitalize()}: {value}")
        print("-" * 80)

    def download_package(self, package_info, download_dir, overwrite=False):
        download_url = f"{self.repo_info.url}/{package_info.filename}"
        file_path = os.path.join(download_dir, os.path.basename(package_info.filename))
        if os.path.exists(file_path):
            if overwrite or self._check_md5(file_path, package_info.md5sum):
                msg = f"{package_info.package}-{package_info.version}:{package_info.architecture} 已存在，{'将直接覆盖下载' if overwrite else '且 MD5 校验通过，跳过下载'}。"
                logging.debug(msg)
                print(msg)
                if not overwrite:
                    print("=" * 80)
                    return True
            else:
                while True:
                    choice = input(
                        f"{package_info.package}-{package_info.version}:{package_info.architecture} 已存在但 MD5 校验不通过，是否覆盖下载？(y/n): ").strip().lower()
                    if choice == 'y':
                        logging.debug(f"用户选择覆盖下载 {package_info.package}-{package_info.version}:{package_info.architecture}。")
                        break
                    elif choice == 'n':
                        logging.debug(f"用户选择跳过下载 {package_info.package}-{package_info.version}:{package_info.architecture}。")
                        print(f"跳过 {package_info.package}-{package_info.version}:{package_info.architecture} 的下载。")
                        print("=" * 80)
                        return True
                    print("输入无效，请输入 'y' 或 'n'。")
        logging.debug(f"开始从 {self.repo_info.name} 下载 {package_info.package}-{package_info.version}:{package_info.architecture}...")
        print(f"开始从 {self.repo_info.name} 下载 {package_info.package}-{package_info.version}:{package_info.architecture}...")
        try:
            response = requests.get(download_url, stream=True)
            total_size = int(response.headers.get('content-length', 0))
            block_size = 1024
            os.makedirs(download_dir, exist_ok=True)
            with open(file_path, 'wb') as file, tqdm(
                    desc=f"{package_info.package}-{package_info.version}",
                    total=total_size,
                    unit='iB',
                    unit_scale=True,
                    unit_divisor=1024
            ) as bar:
                for data in response.iter_content(block_size):
                    bar.update(len(data))
                    file.write(data)
            logging.debug(f"{package_info.package}-{package_info.version}:{package_info.architecture} 已成功从 {self.repo_info.name} 下载到 {download_dir}。")
            print(f"{package_info.package}-{package_info.version}:{package_info.architecture} 已成功从 {self.repo_info.name} 下载到 {download_dir}。")
            self._check_md5(file_path, package_info.md5sum)
            print("=" * 80)
            return True
        except requests.RequestException as e:
            logging.error(f"请求 {download_url} 时出错: {e}")
            print(f"请求 {download_url} 时出错: {e}", file=sys.stderr)
            print(f"下载的文件 {package_info.package}-{package_info.version} 未找到。", file=sys.stderr)
        logging.warning(f"无法从 {self.repo_info.name} 下载 {package_info.package}-{package_info.version}:{package_info.architecture}。")
        print(f"无法从 {self.repo_info.name} 下载 {package_info.package}-{package_info.version}:{package_info.architecture}。")
        with open(os.path.join(download_dir, 'package_missing.log'), 'a') as log_file:
            log_file.write(
                f"{datetime.now()} - 无法从 {self.repo_info.name} 下载 {package_info.package}-{package_info.version}:{package_info.architecture}\n")
        return False

    def _check_md5(self, file_path, expected_md5):
        hash_md5 = hashlib.md5()
        with open(file_path, "rb") as f:
            for chunk in iter(lambda: f.read(4096), b""):
                hash_md5.update(chunk)
        calculated_md5 = hash_md5.hexdigest()
        result = calculated_md5 == expected_md5
        msg = f"{file_path} 的 MD5 校验 {'通过' if result else '不通过，期望: {expected_md5}，实际: {calculated_md5}'}。"
        logging.debug(msg) if result else logging.warning(msg)
        print(msg)
        return result


class DebPackageManager:
    def __init__(self):
        self.repositories = {}

    def add_repository(self, url, architecture, dist, source_types, name=None):
        if name is None:
            name = f"repo_{len(self.repositories) + 1}"
        repo_info = RepositoryInfo(name, url, architecture, dist, source_types)
        self.repositories[name] = Repository(repo_info)

    def query_software_info(self, package_name, fuzzy_search=False):
        found = False
        for repo_name, repo in self.repositories.items():
            print(f"正在 {repo_name} 中查询 {package_name}...")
            if fuzzy_search:
                repo.query_software_info(package_name, fuzzy_search=True)
            else:
                info = repo.get_package_info(package_name)
                if info:
                    found = True
                    repo._print_package_info(info)
        if not found and not fuzzy_search:
            print(f"在所有仓库中均未找到 {package_name} 的软件包信息。")

    def get_dependencies(self, package_info):
        if not package_info.depends:
            return []
        # 移除字符串内所有空格
        dep_string = package_info.depends.replace(" ", "")
        dep_strings = dep_string.split(',')
        dep_infos = []
        for dep in dep_strings:
            alternatives = dep.split('|')
            for alt in alternatives:
                dep_name = re.sub(r'\([^)]*\)', '', alt)
                for repo_name, repo in self.repositories.items():
                    dep_info = repo.get_package_info(dep_name)
                    if dep_info:
                        dep_infos.append(dep_info)
                        break
                if dep_infos and dep_infos[-1].package == dep_name:
                    break
        return dep_infos

    def download_software_and_dependencies(self, package_name, version=None, download_dir=None, download_deps=True,
                                           overwrite=False):
        failed_packages = []
        if download_dir is None:
            download_dir = os.path.join(localsrc_dir, 'main')
        os.makedirs(download_dir, exist_ok=True)
        package_info = None
        for repo_name, repo in self.repositories.items():
            info = repo.get_package_info(package_name)
            if info and (not version or info.version == version):
                package_info = info
                break
        if not package_info:
            logging.warning(f"在所有仓库中均未找到 {package_name}-{version} 的软件包信息，无法下载。")
            print(f"在所有仓库中均未找到 {package_name}-{version} 的软件包信息，无法下载。")
            failed_packages.append(f"{package_name}-{version}")
            with open(os.path.join(download_dir, 'package_missing.log'), 'a') as log_file:
                log_file.write(
                    f"{datetime.now()} - 未找到 {package_name}-{version} 的软件包信息，无法下载。\n")
            if failed_packages:
                print("\n未成功下载的包列表：")
                for pkg in failed_packages:
                    print(f" - {pkg}")
            return

        visited = set()
        stack = [package_info]
        while stack:
            current_package = stack.pop()
            if current_package.package in visited:
                continue
            visited.add(current_package.package)

            for repo_name, repo in self.repositories.items():
                if repo.get_package_info(current_package.package):
                    repo.download_package(current_package, download_dir, overwrite)
                    break

            if download_deps:
                deps = self.get_dependencies(current_package)
                for dep in deps:
                    if dep.package not in visited:
                        stack.append(dep)


    def update_cache(self):
        for repo_name, repo in self.repositories.items():
            repo.update_cache()


if __name__ == "__main__":
    parser = argparse.ArgumentParser(description="查询指定 deb 软件源的软件信息和下载软件及其依赖，固定架构为 arm64。")
    subparsers = parser.add_subparsers(dest='action', help='操作选项')

    # 查询命令
    query_parser = subparsers.add_parser('query', aliases=['q'], help='查询软件包信息')
    query_parser.add_argument('package', help='软件包名称')
    query_parser.add_argument('-f', action='store_true', dest='fuzzy_search', default=False,
                              help='标记启用模糊查询，可在所有包名中进行模糊查询')

    # 下载命令
    download_parser = subparsers.add_parser('download', aliases=['d'], help='下载软件包及其依赖')
    download_parser.add_argument('package', help='软件包名称')
    download_parser.add_argument('version', nargs='?', default=None, help='软件包版本号（可选）')
    download_parser.add_argument('-d', action='store_false', dest='download_deps', default=True,
                                 help='标记仅下载当前包，不下载依赖')
    download_parser.add_argument('-o', dest='download_dir', default=None, help='指定下载目录')
    download_parser.add_argument('-w', action='store_true', dest='overwrite', default=False,
                                 help='标记下载到指定目录时直接覆盖已存在的文件，默认关闭')

    # 更新命令
    update_parser = subparsers.add_parser('update', aliases=['u'], help='更新软件包缓存')

    args = parser.parse_args()

    manager = DebPackageManager()
    # 主仓库
    manager.add_repository("https://archive.kylinos.cn/kylin/KYLIN-ALL", "arm64", "10.1",
                           ["main", "restricted", "universe", "multiverse"], "main_repo")
    # 第三方仓库，修正后的地址
    third_party_url = "https://archive2.kylinos.cn/deb/kylin/production/PART-V10-SP1/custom/partner/V10-SP1"
    manager.add_repository(third_party_url, "arm64", "default", ["all"], "third_party_repo")

    if args.action in ['query', 'q']:
        manager.query_software_info(args.package, args.fuzzy_search)
    elif args.action in ['download', 'd']:
        manager.download_software_and_dependencies(args.package, args.version, args.download_dir,
                                                   download_deps=args.download_deps,
                                                   overwrite=args.overwrite)
    elif args.action in ['update', 'u']:
        manager.update_cache()
    else:
        print("无效的操作。请使用 'query (q)'、'download (d)' 或 'update (u)'，或使用 --help 查看使用说明。", file=sys.stderr)
    